package server

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/textproto"
	"strconv"
	"sync"

	"lsp-server/internal/server/handlers"
)

type Mux struct {
	concurrencyLimit     int64
	requestHandlers      map[string]handlers.Request
	notificationHandlers map[string]handlers.Notification
	writeLock            *sync.Mutex
	conn                 *net.Conn
}

const concurrencyLimit = 4

func NewMux() *Mux {
	return &Mux{
		concurrencyLimit:     concurrencyLimit,
		requestHandlers:      map[string]handlers.Request{},
		notificationHandlers: map[string]handlers.Notification{},
		writeLock:            &sync.Mutex{},
	}
}

func (m *Mux) Process() error {
	for {
		req, err := Read(bufio.NewReader(*m.conn))
		if err != nil {
			return err
		}
		if req.IsNotification() {
			if req.Method != "exit" {
				// Drop notifications sent before initialization.
				slog.Warn("dropping notification sent before initialization", slog.Any("req", req))
				continue
			}
			m.handleMessage(req)
			continue
		}
		if req.Method != "initialize" {
			// Return an error if methods used before initialization.
			slog.Warn("the client sent a method before initialization", slog.Any("req", req))
			if err = m.write(NewResponseError(req.ID, ErrServerNotInitialized)); err != nil {
				return err
			}
			continue
		}
		m.handleMessage(req)
		break
	}
	slog.Info("initialization complete")

	// Handle standard flow.
	sem := make(chan struct{}, m.concurrencyLimit)
	for {
		sem <- struct{}{}
		req, err := Read(bufio.NewReader(*m.conn))
		if err != nil {
			return err
		}
		go func(req Request) {
			m.handleMessage(req)
			<-sem
		}(req)
	}
}

func (m *Mux) handleMessage(req Request) {
	if req.IsNotification() {
		m.handleNotification(req)
		return
	}
	m.handleRequest(req)
}

func (m *Mux) handleNotification(req Request) {
	log := slog.With(slog.String("method", req.Method))
	nh, ok := m.notificationHandlers[req.Method]
	if !ok {
		log.Warn("notification not handled")
		return
	}
	// We don't need to notify clients if the notification results in an error.
	if err := nh.Call(req.Params); err != nil {
		log.Error("failed to handle notification", slog.Any("error", err))
	}
}

func (m *Mux) handleRequest(req Request) {
	log := slog.With(slog.Any("id", req.ID), slog.String("method", req.Method))
	mh, ok := m.requestHandlers[req.Method]
	if !ok {
		log.Error("method not found")
		if err := m.write(NewResponseError(req.ID, ErrMethodNotFound)); err != nil {
			log.Error("failed to respond", slog.Any("error", err))
		}
		return
	}
	var res Response
	result, err := mh.Call(req.Params)
	if err != nil {
		log.Error("failed to handle", slog.Any("error", err))
		res = NewResponseError(req.ID, err)
	} else {
		res = NewResponse(req.ID, result)
	}
	if err = m.write(res); err != nil {
		log.Error("failed to respond", slog.Any("error", err))
	}
}

func (m *Mux) write(msg Message) error {
	m.writeLock.Lock()
	defer m.writeLock.Unlock()
	return Write(bufio.NewWriter(*m.conn), msg)
}

func Read(r *bufio.Reader) (Request, error) {
	// Read header.
	header, err := textproto.NewReader(r).ReadMIMEHeader()
	if err != nil {
		return Request{}, err
	}
	var req Request
	contentLength, err := strconv.ParseInt(header.Get("Content-Length"), 10, 64)
	if err != nil {
		return Request{}, ErrInvalidContentLengthHeader
	}
	// Read body.
	err = json.NewDecoder(io.LimitReader(r, contentLength)).Decode(&req)
	if err != nil {
		return Request{}, err
	}
	if !req.IsJSONRPC() {
		return req, ErrInvalidRequest
	}
	return req, nil
}

func Write(w *bufio.Writer, msg Message) error {
	// Calculate body size.
	body, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	// Write the header.
	_, err = w.WriteString(fmt.Sprintf("Content-Length: %d\r\n\r\n", len(body)))
	if err != nil {
		return err
	}
	// Write the body.
	_, err = w.Write(body)
	if err != nil {
		return err
	}
	// Flush.
	err = w.Flush()
	return err
}
func (m *Mux) Notify(method string, params any) error {
	n := Notification{
		ProtocolVersion: protocolVersion,
		Method:          method,
		Params:          params,
	}
	return m.write(n)
}
func (m *Mux) HandleRequest(name string, request handlers.Request) {
	m.requestHandlers[name] = request
}

func (m *Mux) HandleNotification(name string, notification handlers.Notification) {
	m.notificationHandlers[name] = notification
}

func (m *Mux) NewConn(conn *net.Conn) {
	m.conn = conn
}
